import torch
import torch.nn.functional as F
from PIL import Image
import numpy as np
import cv2

# 将图片转换为张量
# img = Image.open('image.jpg').convert('L')
# tensor = torch.from_numpy(np.array(img)).float() / 255.0
# print(tensor.shape)
# import ipdb
# ipdb.set_trace()
# 使用 Sobel 算子计算梯度
# grad_x = F.conv2d(tensor, torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]), padding=1)
# grad_y = F.conv2d(tensor, torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]]), padding=1)
import cv2
import numpy as np
import os
img_path = "/raid/datasets/cxq_data/MSD/train/mask"
img_list = os.listdir(img_path)
# import ipdb
# ipdb.set_trace()
for img in img_list:
    img_per_name = os.path.join(img_path,img)
    
    # 读取图像
    image = cv2.imread(img_per_name, cv2.IMREAD_GRAYSCALE)

    # 计算水平方向上的梯度
    sobelx = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=3)

    # 计算垂直方向上的梯度
    sobely = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=3)

    # 计算总梯度
    gradient = np.sqrt(sobelx**2 + sobely**2)

    # # 显示结果
    # cv2.imshow('Original Image', image)
    # cv2.imshow('Sobel X', sobelx)
    # cv2.imshow('Sobel Y', sobely)
    # cv2.imshow('Gradient', gradient)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()


    # 将结果转换为 uint8 类型并保存为图片文件
    # edge_map = gradient.clamp(min=0).squeeze().detach().cpu().numpy() * 255
    gradient[gradient!=0]=1
    edge_map = gradient *255
    edge_map = edge_map.astype(np.uint8)
    # import ipdb
    # ipdb.set_trace()
    Image.fromarray(edge_map).save(os.path.join("/raid/datasets/cxq_data/MSD/train/edge",img))
